import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation


class Board(object):
    """Board class that saves all cells.
    """
    def __init__(self, init, t=0):
        """
        :param init: 2D initial board
        :param t: evolution t
        """
        board_ = np.array(init)
        row_n, col_n = board_.shape
        # padding around the board
        self.board_pad = np.zeros((row_n + 2, col_n + 2))
        self.board_pad[1:-1, 1:-1] = board_
        self.board = self.board_pad[1:-1, 1:-1]
        # assert self.board.dim == 2
        self.t = t

    def update_board(self):
        """
        update the board.
        """
        # up down left right upleft upright downleft downright
        score = self.board_pad[0:-2, 1:-1] + \
            self.board_pad[2:, 1:-1] + \
            self.board_pad[1:-1, 0:-2] + \
            self.board_pad[1:-1, 2:] + \
            self.board_pad[0:-2, 0:-2] + \
            self.board_pad[0:-2, 2:] + \
            self.board_pad[2:, 0:-2] + \
            self.board_pad[2:, 2:]

        self.board[(score < 2) | (score > 3) & (self.board == 1)] = 0
        self.board[(score == 3) & (self.board == 0)] = 1

    def evolve(self, plot_flag=False):
        """
        evolution of the board until t.
        """
        if plot_flag:
            fig = plt.figure()
            im = plt.imshow(self.board)
            ax = plt.gca()
            ax.set_title('Iteration: {}'.format(0))
            anim = FuncAnimation(fig, self.plot_board_update, interval=0, fargs=(im, ax))
            plt.show()
        else:
            for _t in   range(self.t):
                self.update_board()

    def plot_board_update(self, *args):
        """plot the board
        """
        # print(args)
        t, im, ax = args
        self.update_board()
        im.set_data(self.board)
        ax.set_title('Iteration: {}'.format(t+1))


board_init = np.random.randint(0, 2, size=(500,500))

# test 4*4 example
# board_init = np.array([[0., 1., 0., 0.],
#                        [0., 0., 1., 0.],
#                        [1., 1., 1., 0.],
#                        [0., 0., 0., 0.]])
b = Board(board_init)
b.evolve(plot_flag=True)
